import matplotlib.pyplot as plt
import numpy as np
from sklearn.neighbors import KNeighborsRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

# 准备数据
x = np.array([182, 178, 170, 168, 165, 162, 158, 154, 149, 144]).reshape(-1, 1)
y = np.array([113, 105, 86, 83, 86, 74, 72, 45, 49, 43])
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.3, random_state=0)

# 计算不同k值的RMSE
k_range = range(2, 8)
rmse_list = []
for k in k_range:
    model = KNeighborsRegressor(n_neighbors=min(k, len(x_train)))
    model.fit(x_train, y_train)
    rmse = np.sqrt(mean_squared_error(y_test, model.predict(x_test)))
    rmse_list.append(rmse)

# 确定最佳k
best_k = list(k_range)[np.argmin(rmse_list)]
best_rmse = np.min(rmse_list)

# 设置中文字体
plt.rcParams.update({
    "font.family": ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"],
    "axes.unicode_minus": False
})

# 绘制RMSE曲线
plt.figure(figsize=(8, 4))
plt.plot(k_range, rmse_list, 'ro-', label='测试集 RMSE')
plt.plot(best_k, best_rmse, 'b*', markersize=15, label=f'最佳 k={best_k}\nRMSE={best_rmse:.2f}')
plt.xlabel('k 的取值')
plt.ylabel('测试集 RMSE')
plt.title('不同 k 值下的预测误差')
plt.grid(alpha=0.3)
plt.legend()
plt.show()

# 绘制拟合曲线
final_model = KNeighborsRegressor(n_neighbors=best_k).fit(x_train, y_train)
x_dense = np.linspace(x.min(), x.max(), 300).reshape(-1, 1)

plt.figure(figsize=(8, 5))
plt.scatter(x, y, s=60, c='k', label='样本点')
plt.plot(x_dense, final_model.predict(x_dense), 'r-', label=f'KNN拟合曲线 (k={best_k})')
plt.xlabel('身高/cm')
plt.ylabel('体重/kg')
plt.title(f'KNN 回归拟合')
plt.axis([140, 190, 40, 140])
plt.grid(alpha=0.3)
plt.legend()
plt.show()